import torch
import argparse

device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")

# -------------------------------------------------------------------------------------------
#	生成torchscript文件
# -------------------------------------------------------------------------------------------
def export_pytorch_model(model_load="./ckpt/checkpoint_iter_120000.pth", input_shape=[256, 256]):
    checkpoint = torch.load(model_load)
    net = torch.load_state_dict(checkpoint["state_dict"])
    net.eval()
    trace_model = torch.jit.trace(net, torch.Tensor(1, 3, input_shape[0], input_shape[1]).to(device))
    print("##"*10)
    print(trace_model)
    print("##" * 10)
    trace_model.save('./ckpt/tmp.pt')

if __name__ == '__main__':
    # 命令行参数
    parser = argparse.ArgumentParser()
    # [model]
    parser.add_argument("--model", help="pth模型", default='./ckpt/checkpoint_iter_120000.pth')
    parser.add_argument("--input_shape", help="模型中的输入尺寸.(图片的resize尺寸)",
                        default='[256,256]')
    args = parser.parse_args()
    export_pytorch_model(args.model, args.input_shape)